# Make an image of a single dataset.

import os
import sys
import matplotlib as mpl
import matplotlib.pyplot as plt
import numpy as np
import readData

# Given a file containing a data set, make some plots.
def plotData(dataFn):
   data = readData.readData(dataFn)

   xPts = []
   yPts = []

   for r in data:
      xPts.append(r[1])
      yPts.append(r[2])

   fig = plt.figure()
   ax = fig.add_subplot(111)

   ax.plot(xPts, yPts, 'ro')
   ax.set_xlim(-4000.0,4000.0)
   ax.set_ylim(-4000.0,4000.0)

   saveFn = dataFn[dataFn.rfind('DD'):]
   plt.savefig("../plots/dataPlot_" + saveFn + ".png")

# Plot the positions of the bodies as a set of orbits.
def plotLines(theta, inc, it):

   xPts = []
   yPts = []

   data = readData.readAllData("../data")

   for r in data:
      for t in range(len(r)):
         p = r[t]
         xPts.append(np.sin(inc)*p[3] + np.cos(inc)*(np.sin(theta)*p[1] + np.cos(theta) * p[2]))
         yPts.append(np.cos(inc)*p[3] + np.sin(inc)*(np.sin(theta)*p[1] + np.cos(theta) * p[2]))


   fig = plt.figure()
   
   ax = fig.add_subplot(111)
   ax.set_xlim(-4000.0,4000.0)
   ax.set_ylim(-4000.0,4000.0)

   ax.plot(xPts, yPts, 'ko', markersize = 1)

   saveName = "../plots/lines_" + str(it).zfill(3) +".png"
   plt.savefig(saveName)

# Plots potential, kinetic, and total energy as a function of time.
def plotEnergy():
   KE = []
   PE = []
   TE = []
   t = []

   data = readData.readAllData("../data")
   outNum = 0

   for DD in data:
     # Compute the kinetic energy.
     k = 0.0

     for r in DD:
        k = k + r[0] * (pow(r[4],2.0) + pow(r[5],2.0) + pow(r[6],2.0)) * 0.5

     KE.append(k)

     # Compute the potential energy
     u = 0.0
     G = 1.0
     for i in range(len(DD)):
        for j in range(i+1,len(DD)):
	   dist = np.sqrt(pow(DD[i][1] - DD[j][1], 2.0)
	         + pow(DD[i][2] - DD[j][2], 2.0)
	         + pow(DD[i][3] - DD[j][3], 2.0));

	   u = u + G * DD[i][0] * DD[j][0] / dist;

     PE.append(u);
     TE.append(k-u);

     t.append(outNum);
     outNum = outNum + 1

   fig = plt.figure()
   ax = fig.add_subplot(111)

   ax.plot(t, KE, 'blue', linewidth = 1, label = "Kinetic")
   ax.plot(t, PE, 'green', linewidth = 1, label = "Potential")
   ax.plot(t, TE, 'black', linewidth= 2, label = "Total")

   plt.savefig("energy.png")
